#!/usr/bin/env python
# -*- coding: utf-8 -*-

from multiprocessing import Process, Queue
import unittest
import time
from Queue import Empty

## class ProcessPool
# responsible for process worker pool management
class ProcessPool(object):
    ## constructor
    # @param num_process: process number in the pool
    # @param in_queue: in message queue
    # @param out_queue: out message queue
    # @param function: function to call 
    def __init__(self, process_name, function):
        self.__process_name = process_name
        self.__num_process = 1
        self._in_queue = None
        self._out_queue = None
        self.__pool_list = []
        self.__function = function
       
    
    ## set the number of process to create
    # @param num_process: number of child process
    def set_num_process(self, num_process):
        self.__num_process = num_process
    
    ## set the incoming queue
    # @param in_queue: the incoming message queue
    def set_incoming_message_queue(self, in_queue):
        self._in_queue = in_queue
        
    ## set the outgoing queue
    # @param _out_queue: the outgoing message queue
    def set_outgoing_message_queue(self, out_queue):
        self._out_queue = out_queue        
    
    ## start all process in pool
    def start_pool(self):
        arguments = ()
        if self._in_queue and self._out_queue:
            arguments = (self._in_queue, self._out_queue,)
        elif self._in_queue:
            arguments = (self._in_queue,)
        elif self._out_queue:
            arguments = (self._out_queue,)
        
        for i in range(self.__num_process):
            a_process = Process(target=self.__function, args=arguments, name=self.__process_name + "-process-"+str(i))
            a_process.daemon = True
            self.__pool_list.append(a_process)
        for a_process in self.__pool_list:
            a_process.start()
            
    ## stop the pool of process
    def stop_pool(self):
        for a_process in self.__pool_list:
            a_process.terminate()
        
        for a_process in self.__pool_list:
            a_process.join()
        
        # reset the process list
        self.__pool_list = []
            
    ## get the process count in the pool
    def get_process_count(self):
        return len(self.__pool_list)
    
    ## get alive process count
    def get_alive_process_count(self):
        count = 0
        for a_process in self.__pool_list:
            if (a_process.is_alive()):
                count = count + 1                
        return count
    
def function_test(in_queue, out_queue):
    while True:
        try:
            s_message = in_queue.get_nowait()
            if s_message:
                out_queue.put_nowait(s_message)              
        except Empty:
            pass
        time.sleep(2)
        
class TestProcessPool(unittest.TestCase):
    def test_init_pool(self):
        """Process pool initialization"""
        a_processpool = ProcessPool("myprocess", function_test)
        a_processpool.set_num_process(5)
        in_queue = Queue()
        out_queue = Queue()
        a_processpool.set_incoming_message_queue(in_queue)
        a_processpool.set_outgoing_message_queue(out_queue)
        
        a_processpool.start_pool()
        self.assertEqual(a_processpool.get_process_count(), 5)
        self.assertEqual(a_processpool.get_alive_process_count(), 5)
        
        a_processpool.stop_pool()
        self.assertEqual(a_processpool.get_process_count(), 0)
        self.assertEqual(a_processpool.get_alive_process_count(), 0)
    
    def test_treatment(self):
        """Test process treatment"""
        a_processpool = ProcessPool("myprocess", function_test)
        a_processpool.set_num_process(10)
        in_queue = Queue()
        out_queue = Queue()
        a_processpool.set_incoming_message_queue(in_queue)
        a_processpool.set_outgoing_message_queue(out_queue)
        
        a_processpool.start_pool()
        
        for i in range(50):
            in_queue.put_nowait("message_" + str(i))
            
        while not in_queue.empty():
            print "Waiting end of treatment..."
            time.sleep(5)            
                
        # check results... check all in messages are in out queue
        result_list = []
        while not out_queue.empty():
            result_list.append(out_queue.get_nowait())
        
        for i in range(50):
            test_message = "message_" + str(i)
            self.assertTrue(test_message in result_list)
        
        
if __name__ == "__main__":
    unittest.main()
